package com.learn.datasource.custom;

import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;

/**
 * @create: 2023-04-18 10:18
 * @author: Mr.Du
 * --------------
 * @notes: 自定义Mysql数据源，从MySQL实时加载数据
 **/
public class CustomMysqlSourceDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

        DataStreamSource<UserInfo> mysqlSource = env.addSource(new MyMysqlRichSource());

        mysqlSource.print();

        env.execute("");
    }
}



class MyMysqlRichSource extends RichSourceFunction<UserInfo> {
    /**定义数据库连接对象*/
    private Connection connection = null;
    /**定义PreparedStatement对象*/
    private PreparedStatement ps = null;

    private boolean isRun = true;

    /**
     * open方法只在实例化类的时候执行一次，所以比较适合用来做数据库连接
     * @param parameters
     * @throws Exception
     */
    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        //加载类驱动
        Class.forName("com.mysql.jdbc.Driver");
        //创建数据库连接
        // 创建数据库连接
        String url = "jdbc:mysql://localhost:3306/test?useUnicode=true&characterEncoding=utf-8&useSSL=false";
        this.connection = DriverManager.getConnection(url, "root", "123456");

        // 准备PreparedStatement对象
        this.ps = connection.prepareStatement("SELECT id, username, password, name FROM user");
    }

    @Override
    public void run(SourceContext<UserInfo> ctx) throws Exception {
        while(isRun) {
            ResultSet resultSet = ps.executeQuery();
            while (resultSet.next()) {
                int id = resultSet.getInt("id");
                String username = resultSet.getString("username");
                String password = resultSet.getString("password");
                String name = resultSet.getString("name");

                ctx.collect(new UserInfo(id, username, password, name));
            }
        }
    }

    /**
    * 使用close方法, 这个方法在销毁实例的时候会执行一次, 比较适合用来关闭连接
    */
    @Override
    public void close() throws Exception {
        super.close();

        // 关闭资源
        if (this.ps != null) {
            this.ps.close();
        }
        if (this.connection != null) {
            this.connection.close();
        }
    }

    @Override
    public void cancel() {
        isRun = false;
        System.out.println("任务被取消......");
    }
}

/**
 数据定义类, POJO
 */
@Data
@AllArgsConstructor
@NoArgsConstructor
class UserInfo {
    int id;
    String username;
    String password;
    String name;
}
